import collections

from sortedcontainers import SortedList


class MKAverage:

    def __init__(self, m: int, k: int):
        self._m, self._k = m, k
        self._queue = collections.deque()  # 数据流中的数值列表
        self._smaller = SortedList()  # 最小的k个值
        self._middle = SortedList()  # 中间用于计算的k个值
        self._bigger = SortedList()  # 最大的k个值
        self._sum = 0  # 中间部分容器的和

    def addElement(self, num1: int) -> None:
        self._queue.append(num1)

        # 当前数据流中的整数还少于m个的情况
        if len(self._queue) < self._m:
            return

        # 当前数据流中的整数正好到达m个的情况
        elif len(self._queue) == self._m:
            queue = sorted(self._queue)
            for i in range(self._k):
                self._smaller.add(queue[i])
            for i in range(self._k, self._m - self._k):
                self._middle.add(queue[i])
            for i in range(self._m - self._k, self._m):
                self._bigger.add(queue[i])
            self._sum = sum(self._middle)

        # 当前数据流中的整数达到m个的情况
        else:
            # 添加新元素
            self._middle.add(num1)
            self._sum += num1

            # 维护中间部分的左右边缘
            self._check_l_and_m()
            self._check_m_and_r()

            # 移除旧元素
            num2 = self._queue.popleft()
            if num2 in self._smaller:
                self._smaller.remove(num2)
                self._change_m_to_l()
            elif num2 in self._bigger:
                self._bigger.remove(num2)
                self._change_m_to_r()
            else:
                self._middle.remove(num2)
                self._sum -= num2

        # print(self._smaller, self._middle, self._bigger)

    def _change_l_to_m(self):
        num = self._smaller.pop()
        self._middle.add(num)
        self._sum += num

    def _change_m_to_r(self):
        num = self._middle.pop()
        self._bigger.add(num)
        self._sum -= num

    def _change_m_to_l(self):
        num = self._middle.pop(0)
        self._smaller.add(num)
        self._sum -= num

    def _change_r_to_m(self):
        num = self._bigger.pop(0)
        self._middle.add(num)
        self._sum += num

    def _check_l_and_m(self):
        if self._smaller[-1] > self._middle[0]:
            num1, num2 = self._smaller.pop(), self._middle.pop(0)
            self._smaller.add(num2)
            self._middle.add(num1)
            self._sum -= num2
            self._sum += num1

    def _check_m_and_r(self):
        if self._middle[-1] > self._bigger[0]:
            num1, num2 = self._middle.pop(), self._bigger.pop(0)
            self._middle.add(num2)
            self._bigger.add(num1)
            self._sum -= num1
            self._sum += num2

    def calculateMKAverage(self) -> int:
        return int(self._sum / (self._m - 2 * self._k)) if len(self._queue) == self._m else -1


if __name__ == "__main__":
    obj = MKAverage(3, 1)
    obj.addElement(3)
    obj.addElement(1)
    print(obj.calculateMKAverage())  # -1
    obj.addElement(10)
    print(obj.calculateMKAverage())  # 3
    obj.addElement(5)
    obj.addElement(5)
    obj.addElement(5)
    print(obj.calculateMKAverage())  # 5
